/* Copyright (c) 2021, National University of Defense Technology. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * 
 *     http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "parcsr_sparsification.h"

namespace YHAMG
{

#define ABS(x) ((x) > 0 ? (x) : -(x))

void ParCSR_Sparsification(const ParCSR_Matrix &A, ParCSR_Matrix &B, double threshold)
{
	MPI_Comm comm = A.comm;
	int n = A.LocLoc.size[0];
	int* A_LocLoc_rowptr = A.LocLoc.rowptr;
	int* A_LocLoc_colind = A.LocLoc.colind;
	double* A_LocLoc_values = A.LocLoc.values;

	int* A_LocExt_rowptr = A.LocExt.rowptr;
	int* A_LocExt_colind = A.LocExt.colind;
	double* A_LocExt_values = A.LocExt.values;

	int A_num_neighbors = A.num_neighbors;
	int* A_neighbors = A.neighbor_ranks;
	int* A_Ext_ptr = A.Ext_ptr;
	int* A_Ext_map = A.Ext_map;

	bool* A_LocLoc_s = new bool[A.LocLoc.rowptr[n]];
	bool* A_LocExt_s = new bool[A.LocExt.rowptr[n]];

	int* B_LocLoc_rowptr = new int[n + 1];
	int* B_LocExt_rowptr = new int[n + 1];

	B_LocLoc_rowptr[0] = 0;
	B_LocExt_rowptr[0] = 0;
#ifdef USE_OPENMP
#pragma omp parallel for schedule(guided)
#endif
	for (int i = 0; i < n; ++i)
	{
		double theta = 0.0;

		for (int j = A_LocLoc_rowptr[i]; j < A_LocLoc_rowptr[i + 1]; ++j)
		{
			if (A_LocLoc_colind[j] == i) continue;
			double abs_jval = ABS(A_LocLoc_values[j]);
			if (abs_jval > theta) theta = abs_jval;
		}

		for (int j = A_LocExt_rowptr[i]; j < A_LocExt_rowptr[i + 1]; ++j)
		{
			double abs_jval = ABS(A_LocExt_values[j]);
			if (abs_jval > theta) theta = abs_jval;
		}

		theta *= threshold;

		int cnt_loc = 1;
		int cnt_ext = 0;

		for (int j = A_LocLoc_rowptr[i]; j < A_LocLoc_rowptr[i + 1]; ++j)
			if (A_LocLoc_s[j] = A_LocLoc_colind[j] != i && ABS(A_LocLoc_values[j]) > theta) ++cnt_loc;

		for (int j = A_LocExt_rowptr[i]; j < A_LocExt_rowptr[i + 1]; ++j)
			if (A_LocExt_s[j] = ABS(A_LocExt_values[j]) > theta) ++cnt_ext;

		B_LocLoc_rowptr[i + 1] = cnt_loc;
		B_LocExt_rowptr[i + 1] = cnt_ext;
	}

	for (int i = 0; i < n; ++i)
		B_LocLoc_rowptr[i + 1] += B_LocLoc_rowptr[i];
	for (int i = 0; i < n; ++i)
		B_LocExt_rowptr[i + 1] += B_LocExt_rowptr[i];

	int* B_LocLoc_colind = new int[B_LocLoc_rowptr[n]];
	double* B_LocLoc_values = new double[B_LocLoc_rowptr[n]];
	int* B_LocExt_colind = new int[B_LocExt_rowptr[n]];
	double* B_LocExt_values = new double[B_LocExt_rowptr[n]];
#ifdef USE_OPENMP
#pragma omp parallel for schedule(guided)
#endif
	for (int i = 0; i < n; ++i)
	{
		double diagonal = 0.0;
		for (int j = A_LocLoc_rowptr[i], k = B_LocLoc_rowptr[i] + 1; j < A_LocLoc_rowptr[i + 1]; ++j)
		{
			if (A_LocLoc_s[j])
			{
				B_LocLoc_colind[k] = A_LocLoc_colind[j];
				B_LocLoc_values[k++] = A_LocLoc_values[j];
			}
			else
				diagonal += A_LocLoc_values[j];
		}

		for (int j = A_LocExt_rowptr[i], k = B_LocExt_rowptr[i]; j < A_LocExt_rowptr[i + 1]; ++j)
		{
			if (A_LocExt_s[j])
			{
				B_LocExt_colind[k] = A_LocExt_colind[j];
				B_LocExt_values[k++] = A_LocExt_values[j];
			}
			else
				diagonal += A_LocExt_values[j];
		}

		B_LocLoc_colind[B_LocLoc_rowptr[i]] = i;
		B_LocLoc_values[B_LocLoc_rowptr[i]] = diagonal;
	}

	delete[] A_LocLoc_s;
	delete[] A_LocExt_s;

	int* B_neighbors = new int[A_num_neighbors];
	for (int r = 0; r < A_num_neighbors; ++r)
		B_neighbors[r] = A_neighbors[r];

	int* v = new int[A.LocExt.size[1]];
	for (int i = 0; i < A.LocExt.size[1]; ++i)
		v[i] = -1;

	for (int j = 0; j < B_LocExt_rowptr[n]; ++j)
		v[B_LocExt_colind[j]] = 1;

	int* B_Ext_map = new int[A.LocExt.size[1]];
	int* B_Ext_ptr = new int[A_num_neighbors + 1];

	B_Ext_ptr[0] = 0;
	for (int r = 0, k = 0; r < A_num_neighbors; B_Ext_ptr[++r] = k)
	{
		for (int i = A_Ext_ptr[r]; i < A_Ext_ptr[r + 1]; ++i)
		{
			if (v[i] == 1)
			{
				v[i] = k;
				B_Ext_map[k++] = A_Ext_map[i];
			}
		}
	}

	for (int j = 0; j < B_LocExt_rowptr[n]; ++j)
		B_LocExt_colind[j] = v[B_LocExt_colind[j]];

	delete[] v;

	B.Free();

	B.comm = comm;

	B.LocLoc.size[0] = n;
	B.LocLoc.size[1] = n;
	B.LocLoc.rowptr = B_LocLoc_rowptr;
	B.LocLoc.colind = B_LocLoc_colind;
	B.LocLoc.values = B_LocLoc_values;

	B.LocExt.size[0] = n;
	B.LocExt.size[1] = B_Ext_ptr[A_num_neighbors];
	B.LocExt.rowptr = B_LocExt_rowptr;
	B.LocExt.colind = B_LocExt_colind;
	B.LocExt.values = B_LocExt_values;

	B.num_neighbors = A_num_neighbors;
	B.neighbor_ranks = B_neighbors;
	B.Ext_ptr = B_Ext_ptr;
	B.Ext_map = B_Ext_map;
}

}